[main][test] Refactor the mtp and eagle test case#5326
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the MTP and Eagle speculative decoding test cases, which is a good improvement for maintainability. However, I've identified several areas with code duplication and misuse of context managers. Specifically, there are redundant calls to cleanup functions and del statements for objects managed by with statements. These issues affect code clarity and could lead to subtle bugs. I've provided suggestions to refactor the duplicated code and remove the unnecessary calls.
| print(f"golden: {golden}") | ||
|
|
||
| assert match | ||
| cleanup_dist_env_and_memory() |
| if not enforce_eager: | ||
| if cudagraph_mode == "FULL_DECODE_ONLY": | ||
| pytest.skip("This case will be supported in future") | ||
| with VllmRunner(model_name, | ||
| tensor_parallel_size=4, | ||
| max_model_len=4096, | ||
| gpu_memory_utilization=0.8, | ||
| distributed_executor_backend="mp", | ||
| speculative_config={ | ||
| "method": | ||
| "mtp", | ||
| "num_speculative_tokens": | ||
| num_speculative_tokens, | ||
| "disable_padded_drafter_batch": | ||
| disable_padded_drafter_batch, | ||
| }, | ||
| enforce_eager=enforce_eager, | ||
| compilation_config=CompilationConfig( | ||
| cudagraph_mode=cudagraph_mode, | ||
| cudagraph_capture_sizes=[12], | ||
| )) as spec_llm: | ||
| spec_outputs = spec_llm.generate_greedy(example_prompts, max_tokens) | ||
| del spec_llm | ||
|
|
||
| else: | ||
| if cudagraph_mode == "PIECEWISE": | ||
| pytest.skip("skipping the repeating case") | ||
| with VllmRunner(model_name, | ||
| tensor_parallel_size=4, | ||
| max_model_len=4096, | ||
| gpu_memory_utilization=0.8, | ||
| distributed_executor_backend="mp", | ||
| speculative_config={ | ||
| "method": | ||
| "mtp", | ||
| "num_speculative_tokens": | ||
| num_speculative_tokens, | ||
| "disable_padded_drafter_batch": | ||
| disable_padded_drafter_batch, | ||
| }, | ||
| enforce_eager=enforce_eager, | ||
| ) as spec_llm: | ||
| spec_outputs = spec_llm.generate_greedy(example_prompts, max_tokens) | ||
| del spec_llm |
There was a problem hiding this comment.
The if/else block for enforce_eager contains a lot of duplicated code for initializing VllmRunner. This can be refactored to reduce redundancy and improve readability. Additionally, the del spec_llm calls are unnecessary when using a with statement, as the context manager handles cleanup.
compilation_config = None
if not enforce_eager:
if cudagraph_mode == "FULL_DECODE_ONLY":
pytest.skip("This case will be supported in future")
compilation_config = CompilationConfig(
cudagraph_mode=cudagraph_mode,
cudagraph_capture_sizes=[12],
)
else:
if cudagraph_mode == "PIECEWISE":
pytest.skip("skipping the repeating case")
with VllmRunner(model_name,
tensor_parallel_size=4,
max_model_len=4096,
gpu_memory_utilization=0.8,
distributed_executor_backend="mp",
speculative_config={
"method":
"mtp",
"num_speculative_tokens":
num_speculative_tokens,
"disable_padded_drafter_batch":
disable_padded_drafter_batch,
},
enforce_eager=enforce_eager,
compilation_config=compilation_config) as spec_llm:
spec_outputs = spec_llm.generate_greedy(example_prompts, max_tokens)| # Heuristic: expect at least 66% of the prompts to match exactly | ||
| # Upon failure, inspect the outputs to check for inaccuracy. | ||
| assert matches > int(0.66 * len(ref_outputs)) | ||
| cleanup_dist_env_and_memory() |
| if not enforce_eager: | ||
| with VllmRunner(model_name, | ||
| tensor_parallel_size=1, | ||
| max_num_seqs=256, | ||
| gpu_memory_utilization=0.7, | ||
| distributed_executor_backend="mp", | ||
| enable_expert_parallel=True, | ||
| speculative_config={ | ||
| "method": | ||
| "mtp", | ||
| "num_speculative_tokens": | ||
| num_speculative_tokens, | ||
| "disable_padded_drafter_batch": | ||
| disable_padded_drafter_batch, | ||
| }, | ||
| enforce_eager=enforce_eager, | ||
| max_model_len=2000, | ||
| compilation_config=CompilationConfig( | ||
| cudagraph_mode=cudagraph_mode, | ||
| cudagraph_capture_sizes=[12], | ||
| )) as spec_llm: | ||
| sampling_config = SamplingParams(temperature=0, max_tokens=256, ignore_eos=False) | ||
| spec_outputs = spec_llm.generate(example_prompts, sampling_config) | ||
|
|
||
| else: | ||
| if cudagraph_mode == "PIECEWISE": | ||
| pytest.skip("skipping the repeating case") | ||
| with VllmRunner(model_name, | ||
| tensor_parallel_size=1, | ||
| max_num_seqs=256, | ||
| gpu_memory_utilization=0.7, | ||
| distributed_executor_backend="mp", | ||
| enable_expert_parallel=True, | ||
| speculative_config={ | ||
| "method": | ||
| "mtp", | ||
| "num_speculative_tokens": | ||
| num_speculative_tokens, | ||
| "disable_padded_drafter_batch": | ||
| disable_padded_drafter_batch, | ||
| }, | ||
| enforce_eager=enforce_eager, | ||
| max_model_len=2000 | ||
| ) as spec_llm: | ||
| sampling_config = SamplingParams(temperature=0, max_tokens=256, ignore_eos=False) | ||
| spec_outputs = spec_llm.generate(example_prompts, sampling_config) |
There was a problem hiding this comment.
This if/else block for enforce_eager has a lot of duplicated code for VllmRunner initialization. It can be refactored to be more concise and maintainable.
compilation_config = None
if not enforce_eager:
compilation_config = CompilationConfig(
cudagraph_mode=cudagraph_mode,
cudagraph_capture_sizes=[12],
)
else:
if cudagraph_mode == "PIECEWISE":
pytest.skip("skipping the repeating case")
with VllmRunner(model_name,
tensor_parallel_size=1,
max_num_seqs=256,
gpu_memory_utilization=0.7,
distributed_executor_backend="mp",
enable_expert_parallel=True,
speculative_config={
"method":
"mtp",
"num_speculative_tokens":
num_speculative_tokens,
"disable_padded_drafter_batch":
disable_padded_drafter_batch,
},
enforce_eager=enforce_eager,
max_model_len=2000,
compilation_config=compilation_config) as spec_llm:
sampling_config = SamplingParams(temperature=0, max_tokens=256, ignore_eos=False)
spec_outputs = spec_llm.generate(example_prompts, sampling_config)| cleanup_dist_env_and_memory() | ||
| del spec_llm |
| cleanup_dist_env_and_memory() | ||
| del llm |
| cleanup_dist_env_and_memory() | ||
| del llm |
| VLLM_WORKER_MULTIPROC_METHOD: spawn | ||
| if: ${{ inputs.type == 'full' }} | ||
| run: | | ||
| pytest -sv --durations=0 tests/e2e/multicard/spec_decode_v1/test_mtp_qwen3_next.py |
There was a problem hiding this comment.
move to 4-card test part
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
|
|
||
| @pytest.mark.parametrize("model_name", MODELS) | ||
| @pytest.mark.parametrize("num_speculative_tokens", [1,2,3]) | ||
| @pytest.mark.parametrize("enforce_eager", [True, False]) |
There was a problem hiding this comment.
remove enforce_eager=True
| @pytest.mark.parametrize( | ||
| "cudagraph_mode", | ||
| [ | ||
| CUDAGraphMode.NONE, |
| [ | ||
| CUDAGraphMode.NONE, | ||
| CUDAGraphMode.PIECEWISE, | ||
| CUDAGraphMode.FULL_DECODE_ONLY, |
There was a problem hiding this comment.
add it back once FULL_DECODE_ONLY works
| @pytest.mark.parametrize("method", ["eagle", "eagle3"]) | ||
| @pytest.mark.parametrize("disable_padded_drafter_batch", [True, False]) | ||
| @pytest.mark.parametrize("async_scheduling", [True, False]) | ||
| def test_offline_eagle_correctness(model_name: str, |
There was a problem hiding this comment.
test_<model_name>
test_deepseek_mtp_accuracy
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
|
There are some conflicts, please rebase |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: lilinsiman <lilinsiman@gmail.com>
Signed-off-by: lilinsiman <lilinsiman@gmail.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: lilinsiman <lilinsiman@gmail.com>
…to FIA_rebase * 'main' of https://github.com/vllm-project/vllm-ascend: [feature] mooncake support pcp/dcp in common conditions (vllm-project#5224) [Bugfix] Fix mm_merge (vllm-project#5249) [Main2Main] Upgrade vllm commit to 1230 (vllm-project#5495) [Feature] Refactor PCP &DCP related code (vllm-project#5214) [main][test] Refactor the mtp and eagle test case (vllm-project#5326) [smoke][bugfix] moe_init_routing_v2 active_expert_range use int type (vllm-project#5521) [2/N] Upgrade nightly doc (vllm-project#5534) [Doc] Add new contributors. (vllm-project#5537) [3/N][Nightly] Move ops tests to nightly (vllm-project#5538)
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com>
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com>
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com> Signed-off-by: zrj026 <zhangrunjiang026@gmail.com>
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com>
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com> Signed-off-by: zrj026 <zhangrunjiang026@gmail.com>
### What this PR does / why we need it? 1. Refactor the current test with mtp and eagle cases 2. Add new necessary cases with mtp and eagle ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ut - vLLM version: release/v0.13.0 - vLLM main: vllm-project/vllm@5fbfa8d --------- Signed-off-by: lilinsiman <lilinsiman@gmail.com>
What this PR does / why we need it?
Does this PR introduce any user-facing change?
no
How was this patch tested?
ut